using DataFrames
using QuadGK
using Distributions
using CSV
using LinearAlgebra
using JLD
using Interpolations
using DelimitedFiles

#clearconsole()

#-------------------------------------------------------------
# IRF Configuration
#-------------------------------------------------------------

H = 40 # maximum horizon
sh_size = 3   # shock size, in multiples of standard deviations
n_drawsread = 1000 # transform draws 1 to n_drawsread

#-------------------------------------------------------------
# include Functions
#-------------------------------------------------------------
#cd("$(pwd())/Dropbox/Heterogeneity/Software/KS_Simulation/")
readDir = "$(pwd())/Functions/"
# include(readDir *"vech.jl");
include(readDir *"logSpline_Procedures.jl");
# include(readDir *"VAR_Procedures.jl");
include(readDir *"Loaddata.jl");
include(readDir *"EmpPercentiles_Procedures.jl")
include(readDir *"GiniFracBelowCutoff_Procedures.jl")

#-------------------------------------------------------------
# choose specification files
#-------------------------------------------------------------
nfVARSpec = "10tc"
nModSpec  = "1" # K=10
#nModSpec  = "2" # K=22
nMCMCSpec = "1"
modName   = "SS"  # VAR or SS
cutoff    = 1;
#nShockspec= "AggSh1"
nShockspec= "DistrSh" #"AggSh1"
# Aggregate shocks 1: TFP shock, 2: GDP shock,  3: Employment shock
# Distr shock DistrSh is Gini on impact


specDir   = "$(pwd())/SpecFiles/"
include(specDir * "/fVARspec" * nfVARSpec * ".jl")
include(specDir * "/" * modName * "spec" * nModSpec * ".jl")
include(specDir * "/" * modName * "MCMCspec" * nMCMCSpec * ".jl")

# Percentiles of interest
ngrid     = 20
grid_temp = range(0.2, stop=1.8, length=ngrid);
grid_temp = [0.0; grid_temp]
vec_percs = [0.1; 0.2; 0.5; 0.8; 0.9]

#-------------------------------------------------------------
# load aggregate data
#-------------------------------------------------------------
dataDir  = "$(pwd())/data/"
agg_data, period_agg, mean_unrate = loadaggdata(SampleStart,SampleEnd)
n_agg    = size(agg_data)[2]

#-------------------------------------------------------------
# Load coefficients from density estimation
#-------------------------------------------------------------
sNameLoadDir = "fVAR" * nfVARSpec
loaddir      = "$(pwd())/results/" * sNameLoadDir *"/";
sNameFile    = "K" * string(K) * "_fVAR" * nfVARSpec

PhatDensCoef_factor, MDD_GoF, VinvLam_all, period_Dens_ind, PhatDensCoef_lambda, PhatDensCoef_mean, PhatDensCoef_mean_allt  = loaddensdata(SampleStart,SampleEnd,K,nfVARSpec)
n_cross = size(PhatDensCoef_factor)[2]

knots_all = readdlm(loaddir * "fVAR" *nfVARSpec * "_knots_all.csv", ',', Float64, '\n'; skipstart = 1);
ii=getindex.(findall(K_vec.-K.==0),[1 2])[1] # find index ii where K==K_vec
knots = knots_all[quant_sel[ii,:].==1]

#-------------------------------------------------------------
# Load posterior means
#-------------------------------------------------------------

sName    = "fVAR" * nfVARSpec * "_" * modName * nModSpec * "_" * "MCMC" * nMCMCSpec;
loadDir  = "$(pwd())/results/" * sName *"/";
Bpdraw   = load(loadDir * sName * "_PostDraws.jld", "Bpdraw")
#n_subseq                 = floor(Int,size(PHIpdraw)[1]/n_every)

YY_IRF = CSV.read(loadDir * sName * "_IRF_YY_" *nShockspec * "_pmean.csv", DataFrame, header = true);
PhatDens_IRF = CSV.read(loadDir * sName * "_IRF_PhatDens_" *nShockspec*"_pmean.csv", DataFrame, header = true);
#PhatDens_IRF = CSV.read(loadDir * sName * "_IRF_PhatDens_" *nShockspec*"_pmean_nozeros.csv", DataFrame, header = true);
YY_IRF = convert(Array, YY_IRF)
PhatDens_IRF = convert(Array, PhatDens_IRF)

#-------------------------------------------------------------
# Generate transformed IRFs at posterior mean
#-------------------------------------------------------------

println("")
println("Generating transformed IRFs at posterior mean... ")
time_init_loop = time_ns();

emp_percs_all    = zeros(H+1,length(vec_percs))
densvalues_ss    = PhatDens_IRF[1,:]; # steady state density
PhatMassDiff_IRF = zeros(H+1,1);
Gini_IRF         = zeros(H+1,1);

for hh = 1:(H+1)
    # read estimated states in period t, 1:n_agg are agg data, n_agg+1:end are dens coeff
    statepmean_hh = YY_IRF[hh,:]

    PhatDensCoef_hh    = coefRecover(statepmean_hh[n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
#    emp_percs_all[hh,:] = DensPercentiles(PhatDensCoef_hh, knots, 0.0, xgrid, grid_temp, vec_percs) # not including zeros
    emp_percs_all[hh,:] = DensPercentiles(PhatDensCoef_hh, knots, statepmean_hh[3]+mean_unrate, xgrid, grid_temp, vec_percs)
    PhatMassDiff_IRF[hh,1] = ProbMassDiff_Cutoff(PhatDens_IRF[hh,:], densvalues_ss, cutoff, xgrid);
    Gini_IRF[hh] = GiniCoef(PhatDens_IRF[hh,:], xgrid)

end

savedir = "$(pwd())/results/" * sName *"/";
CSV.write(savedir * sName * "_IRF_Pctl_" *nShockspec*"_pmean.csv", DataFrame(emp_percs_all))
CSV.write(savedir * sName * "_IRF_BelowCutoff_" *nShockspec*"_pmean.csv", DataFrame(PhatMassDiff_IRF))
CSV.write(savedir * sName * "_IRF_Gini_" *nShockspec*"_pmean.csv", DataFrame(Gini_IRF))
#CSV.write(savedir * sName * "_IRF_Pctl_" *nShockspec*"_pmean_nozeros.csv", DataFrame(emp_percs_all))
#CSV.write(savedir * sName * "_IRF_BelowCutoff_" *nShockspec*"_pmean_nozeros.csv", DataFrame(PhatMassDiff_IRF))
#CSV.write(savedir * sName * "_IRF_Gini_" *nShockspec*"_pmean_nozeros.csv", DataFrame(Gini_IRF))

time_loop=signed(time_ns()-time_init_loop)/1000000000
println("Elapsed time = $(time_loop)")
println("")


#-------------------------------------------------------------
# Generate IRFs for a subset of posterior draws
#-------------------------------------------------------------

println("")
println("Generating transformed IRFs for posterior draws... ")
println("")

YY_IRF_uncertainty        = zeros(H+1,n_agg+n_cross,n_drawsread)
PhatDens_IRF_uncertainty  = zeros(H+1,length(xgrid),n_drawsread)

for pp = 1:n_drawsread
    YY_IRF_pp = CSV.read(loadDir * sName * "_IRF_YY_" *nShockspec * "_"*string(pp) *".csv", DataFrame, header = true);
    PhatDens_IRF_pp = CSV.read(loadDir * sName * "_IRF_PhatDens_" *nShockspec* "_" * string(pp) * ".csv", DataFrame, header = true);
#    PhatDens_IRF_pp = CSV.read(loadDir * sName * "_IRF_PhatDens_" *nShockspec* "_" * string(pp) * "_nozeros.csv", DataFrame, header = true);

    YY_IRF_uncertainty[:,:,pp] = convert(Array, YY_IRF_pp)
    PhatDens_IRF_uncertainty[:,:,pp] = convert(Array, PhatDens_IRF_pp)
end

emp_percs_uncertainty        = zeros(H+1,length(vec_percs),n_drawsread)
PhatMassDiff_IRF_uncertainty = zeros(H+1,n_drawsread);
Gini_IRF_uncertainty         = zeros(H+1,n_drawsread);

for pp = 1:n_drawsread

    time_init_loop = time_ns();
    println("Draw number:  $pp")
    println("Remaining draws:  $(n_drawsread-pp)")

    for hh = 1:(H+1)
        # read estimated states in period t, 1:n_agg are agg data, n_agg+1:end are dens coeff
        statepmean_hh = YY_IRF_uncertainty[hh,:,pp]

        PhatDensCoef_hh    = coefRecover(statepmean_hh[n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
#        emp_percs_uncertainty[hh,:,pp] = DensPercentiles(PhatDensCoef_hh, knots, 0.0, xgrid, grid_temp, vec_percs) # no zeros
        emp_percs_uncertainty[hh,:,pp] = DensPercentiles(PhatDensCoef_hh, knots, statepmean_hh[3]+mean_unrate, xgrid, grid_temp, vec_percs)
        PhatMassDiff_IRF_uncertainty[hh,pp] = ProbMassDiff_Cutoff(PhatDens_IRF_uncertainty[hh,:,pp], densvalues_ss, cutoff, xgrid);
        Gini_IRF_uncertainty[hh,pp] = GiniCoef(PhatDens_IRF_uncertainty[hh,:,pp], xgrid)
    end

    time_loop=signed(time_ns()-time_init_loop)/1000000000
    println("Elapsed time = $(time_loop)")
    println("")
    CSV.write(savedir * sName * "_IRF_Pctl_" *nShockspec * "_" * string(pp)* ".csv", DataFrame(emp_percs_uncertainty[:,:,pp]))
#    CSV.write(savedir * sName * "_IRF_Pctl_" *nShockspec * "_" * string(pp)* "_nozeros.csv", DataFrame(emp_percs_uncertainty[:,:,pp]))

end

CSV.write(savedir * sName * "_IRF_BelowCutoff_" *nShockspec*"_uncertainty.csv", DataFrame(PhatMassDiff_IRF_uncertainty))
CSV.write(savedir * sName * "_IRF_Gini_" *nShockspec*"_uncertainty.csv", DataFrame(Gini_IRF_uncertainty))
#CSV.write(savedir * sName * "_IRF_BelowCutoff_" *nShockspec*"_uncertainty_nozeros.csv", DataFrame(PhatMassDiff_IRF_uncertainty))
#CSV.write(savedir * sName * "_IRF_Gini_" *nShockspec*"_uncertainty_nozeros.csv", DataFrame(Gini_IRF_uncertainty))
